from typing import Optional

import requests
from fastapi import HTTPException
from pydantic import BaseModel

import config


# one api后台操作认证
headers = {
    "Authorization": f"Bearer {config.model['backend_token']}"
}


class ModelInfo(BaseModel):
    base_url: str = ""
    config: str = "{}"
    group: str = "default"
    groups: list[str] = ["default"]
    is_edit: bool = False
    key: str = "EMPTY"
    model_mapping: str = ""
    models: str
    name: str
    other: str = ""
    type: int


class UpdateModelInfo(BaseModel):
    base_url: str = ""
    config: str = "{}"
    group: str = "default"
    groups: list[str] = ["default"]
    is_edit: bool = True
    key: str = ""
    model_mapping: str = ""
    models: Optional[str] = ""
    name: str
    other: str = ""
    type: int
    balance: int = 0
    balance_updated_time: int =0
    id: int
    status: int = 1
    system_prompt: str = None



def get_models(model_type: str):
    # 访问one api接口获取模型列表
    res = []
    _models = []
    i = 0
    # 分页获取所有模型信息
    while True:
        response = requests.get(config.model["backend_url"], params={"p": i}, headers=headers)
        if response.status_code == 200:
            response_json = response.json()
            if len(response_json["data"]) == 0:
                break
            _models += response_json["data"]
            i += 1
        else:
            raise HTTPException(detail=response.text, status_code=response.status_code)
    # 根据模型类型筛选
    if model_type == "rerank":
        for _model in _models:
            _sub_models = _model["models"].split(",")
            if _model["type"] != 43:
                continue
            res.append({"type": _model["name"], "models": _sub_models})
    else:
        for _model in _models:
            if _model["type"] == 43:
                continue
            _sub_models = _model["models"].split(",")
            if model_type == "llm":
                models = [item for item in _sub_models if "embed" not in item]
            else:
                models = [item for item in _sub_models if "embed" in item]
            if len(models) == 0:
                continue
            res.append({"type": config.MODEL_OPTIONS[str(_model["type"])], "models": models})

    return {"code": 200, "data": res, "message": f"获取{model_type}模型列表成功"}
